import numpy as np
import json

class RBAC(object):
    """
    =======================================================================================================
    Python Implementation for RBAC System. You can easily use it with some configs.
    ------------------------------
    API:
    1. return_actions;
    2. save_json;
    3. loadJsonWithMatrix;
    4. loadAllJsonConfig;
    5. addUser;
    6. addAction;
    7. addRole;
    8. removeUser;
    9. removeAction;
    10. removeRole;
    ------------------------------
    Json Config File Example:
        {"users": ["id001", "id002", "id003"],
          "roles": ["admin", "lv3", "lv2"],
          "actions": ["money10000", "moneyPlus1", "automaticPlus1"],
          "cannot_actions": ["say_cannot", "say_cannot", "say_cannot"],
          "user_role_matrix": [[1.0, 0.0, 0.0], [0.0, 1.0, 0.0], [0.0, 0.0, 1.0]],
          "role_action_matrix": [[1.0, 1.0, 1.0], [0.0, 1.0, 1.0], [0.0, 0.0, 1.0]],
          "user_roles_dict": {
            "id001": ["admin"],
            "id002": ["lv3"],
            "id003": ["lv2"]},
          "action_roles_dict":
          {"money10000": ["admin"],
            "moneyPlus1": ["admin", "lv3"],
            "automaticPlus1": ["admin", "lv3", "lv2"]}}
    ------------------------------
    Running Test Example:

    IN test.py:

        from rbac import RBAC
        users=["id001", "id002", "id003"]
        roles=["admin", "lv3", "lv2"]
        actions=["money10000","moneyPlus1","automaticPlus1"]
        cannot_actions=["say_cannot"]*len(actions)

        user_role_matrix=[[1,0,0],
                          [0,1,0],
                          [0,0,1]]

        role_action_matrix=[[1,1,1],
                            [0,1,1],
                            [0,0,1]]

        my_RBAC=RBAC(users,roles,actions,cannot_actions,user_role_matrix,role_action_matrix)
        print(my_RBAC.return_actions("id003","money10000"))
        my_RBAC.save_json("config.json")

        my_RBAC.loadAllJsonConfig("config.json")
    ------------------------------
    LiangZi, XJTU, 05, Jun, 2021
    UNDER MIT 2.0 LICENSE
    =======================================================================================================
    """
    def __init__(self, user_list=[], role_list=[], action_list=[], cannot_action_list=[], user_role_matrix=[[]], role_action_matrix=[[]]):
        """init function
        ---------------------------------------------------------------------------------------------------
        user_list: user list
        role_list: role list
        action_list: action list
        cannot_action_list: action list when have no permission to use action in action_list.
        user_role_matrix: 0-1 matrix, elements_ij in which indicates user[i] have the role of role[j]
        role_action_matrix: 0-1matrix, elements_ij in which indicates role[i] have the action of action[j]
        ---------------------------------------------------------------------------------------------------
        """
        super(RBAC, self).__init__()

        ## data save
        self.user_list=user_list
        self.role_list=role_list
        self.action_list=action_list
        self.cannot_action_list=cannot_action_list
        
        self.user_role_matrix=np.array(user_role_matrix,dtype=float)
        self.role_action_matrix=np.array(role_action_matrix,dtype=float)

        #meta_data
        self.num_user=None
        self.num_role=None
        self.num_action=None

        # generated data map
        self.user_roles_dict=None
        self.action_roles_dict=None
        self.__init_for_map()

    def __init_for_map(self):
        """construct two map dict from lists and matrices."""
       
        if self.user_roles_dict is not None:
            print("[[[[[SKIP]]]]]")
            return 0

        self.reloadForMap()
        
    def reloadForMap(self):
        """construct two map dict from lists and matrices."""
        
        print(">>>Begin to construct map from user to roles.")
        self.user_roles_dict={}
        for index, user in enumerate(self.user_list):
            self.user_roles_dict[user]=[]
            index_role_list=self.user_role_matrix[index]
            for jndex, ele in enumerate(index_role_list):
                if ele-0==0.:
                    continue
                else:
                    self.user_roles_dict[user].append(self.role_list[jndex])
                
        print(">>>Begin to construct map from role to actions.")
        self.action_roles_dict={}
        for index, action in enumerate(self.action_list):
            self.action_roles_dict[action]=[]
            
            index_roles_list=(self.role_action_matrix.T)[index]
            for jndex,ele in enumerate(index_roles_list):
                if ele-0==0.:
                    continue
                else:
                    self.action_roles_dict[action].append((self.role_list[jndex]))
        print("Map Construction DONE.\n")
        
        
    def return_actions(self, user, action):
        """if `user` have the permission of `action`, then return `action`, else use `cannot action`. """
        if self.user_roles_dict==None:
            self.__init_for_map()

        print(self.user_roles_dict)
        print(self.action_roles_dict)

        roles1=self.user_roles_dict[user]
        # print(roles1)
        roles2=self.action_roles_dict[action]
        # print(ro)
        print("===="*4+"inference"+"===="*4)
        myaction=None
        if len(set(roles1)& set(roles2))!=0:
            # print("===="*4+"inference"+"===="*4)
            print("We allow [{}] becuase of roles [{}] for action [{}]".format(user,(set(roles1)&set(roles2)),action))
            myaction=action
        else:
            print("CANNOT allow [{}] for action [{}], we replaced it with the action [{}]".format(user,action,self.cannot_action_list[self.action_list.index(action)]))
            myaction= self.cannot_action_list[self.action_list.index(action)]
        print("===="*2+"END-inference"+"===="*4)
        return myaction

    def save_json(self,path="rbac_data.json"):
        """save config file into `path`. """
        action_list=[]
        cannot_action_list=[]
        action_roles_dict={}
        for action in self.action_list:
            action_list.append(action.__name__)
        for action in self.cannot_action_list:
            cannot_action_list.append(action.__name__)
        for key in self.action_roles_dict:
            action_roles_dict[key.__name__]=self.action_roles_dict[key]
        bigdict={"users":self.user_list,
                 "roles":self.role_list,
                 "actions":action_list,
                 "cannot_actions":cannot_action_list,
                 "user_role_matrix":self.array2lists(self.user_role_matrix),
                 "role_action_matrix":self.array2lists(self.role_action_matrix),
                 "user_roles_dict":self.user_roles_dict,
                 "action_roles_dict":action_roles_dict,
                 }
        print(bigdict)
        with open(path,'w') as f:
            json.dump(bigdict,f)
            
        print("---->json file save done.")

    def loadJsonWithMatrix(self,load_path):
        """only load matrix information from `load_path`"""
        
        with open(load_path,'r') as f:
            bigdict=json.load(f)

        self.__init__(bigdict["users"],bigdict["roles"],bigdict["actions"],
                    bigdict["user_role_matrix"],bigdict["role_action_matrix"])
        self.__init_for_map()


    def loadAllJsonConfig(self,load_path):
        """load for initing."""
        
        with open(load_path,'r') as f:
            bigdict=json.load(f)

        self.__init__(bigdict["users"],bigdict["roles"],bigdict["actions"],bigdict["cannot_actions"],
                      bigdict["user_role_matrix"],bigdict["role_action_matrix"])

        
    def array2lists(self, array):
        """2darray to lists"""
        
        ls=[]
        
        for ele in array:
            ls.append(list(ele))
        return ls


    def addUser(self,user,user_roles_list=None, user_roles_dict=None):
        """add `user` to RBAC system, with user-role-list, if not, use user-roles-dict"""

        self.user_list.append(user)
        user_role_lss=self.array2lists(self.user_role_matrix)
        if user_roles_list is not None:
            user_role_lss.append(user_role_lss)
        else:
            user_role_lss.append(user_roles_dict[user])
        self.user_role_matrix=np.array(user_role_lss)

        self.reloadForMap()

    def addAction(self, action,cannot_action_list, action_roles_list=None, action_roles_dict=None):
        """add `action` to RBAC system, with action-roles-list, if not, use action-roles-dict"""
        
        
        self.action_list.append(action)
        self.cannot_action_list.append(cannot_action)
        
        action_role_lss=self.array2lists(self.role_action_matrix.T)
        if action_roles_list is not None:
            action_role_lss.append(action_role_lss)
        else:
            action_role_lss.append(action_roles_dict[action])
        self.role_action_matrix=np.array(action_role_lss).T
        self.reloadForMap()

    def addRole(self, role, role_users_list, role_action_list):
        """add `role` to RBAC system, with role-users-list, if not, use role-action-dict"""

        ## add role
        self.role_list.append(role)

        ## update user role matrix
        user_role_lss=self.array2lists(self.user_role_matrix)
        user_role_lss.append(role_users_list)
        self.user_role_matrix=np.array(user_role_lss)

        # updat role action matrix
        self.role_action_matrix=np.array(self.array2lists(self.role_action_matrix).append(role_action_list))

        self.reloadForMap()

        
    def removeUser(self,user):
        """remove `user` from RBAC system"""
        
        index=self.user_list.index(user)

        self.user_list.remove(user)
        self.user_role_matrix=np.delete(self.user_role_matrix,index,axis=0)
        self.reloadForMap()

    def removeAction(self,action):
        """remove `action` from RBAC system"""
      
        index=self.action_list.index(user)

        self.action_list.remove(user)
        self.role_action_matrix=np.delete(self.role_action_matrix,index,axis=1)
        self.reloadForMap()

    def removeRole(self, role):
        """remove `role` from RBAC system"""
      
        index=self.role_list.index(role)

        self.role_list.remove(role)
        self.user_role_matrix=np.delete(self.user_role_matrix,index,axis=1)
        self.role_action_matrix=np.delete(self.role_action_matrix,index,axis=0)
        self.reloadForMap()
